import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker

# 带宽条件，单位为 Gbps
bandwidths = [1, 10, 25, 50]  # 带宽条件

# 不同模型和方法在不同带宽条件下的性能数据 (速度或吞吐量)
baseline_vgg = [1.53, 22.31, 32.93, 36.45]
topk_vgg = [8.3, 25.94, 26.91, 27.63]
netsense_vgg = [180.88, 184.88, 190.88]

bandwidths_resnet = [1, 10, 25, 50]

baseline_resnet = [49.81, 55.38, 62.97, 112.60]
topk_resnet = [119.78, 120.78, 122.54, 123.03]
netsense_resnet = [121.27, 123.27, 125.29, 127.60]

baseline_bert = [12, 25, 53.13]
topk_bert = [39.13, 40.13, 54.13]
netsense_bert = [53.13, 54.13, 55.13]

# 设置柱状图的宽度
bar_width = 0.25
index = np.arange(len(bandwidths))  # 确保 index 长度与数据长度匹配

# 设置配色和填充模式
colors = ['#B4D7E9', '#5CB95D', '#FF9300']  # 自定义颜色
hatch_patterns = ['//', '\\\\', '--']  # 斜线、十字、实心填充

# 创建图像和子图 (1x3 布局)
fig, axs = plt.subplots(1, 3, figsize=(16, 4))  # 1行3列布局

# 设置统一的柱状图样式
def plot_bars(ax, baseline, topk, netsense, title, ylabel, ylim):
    ax.bar(index, baseline, bar_width, label='Baseline', color=colors[0], hatch=hatch_patterns[0])
    ax.bar(index + bar_width, topk, bar_width, label='TopK 0.1', color=colors[1], hatch=hatch_patterns[1])
    ax.bar(index + 2 * bar_width, netsense, bar_width, label='NetSenseML', color=colors[2], hatch=hatch_patterns[2])
    ax.set_xlabel('Bandwidth (Gbps)')
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xticks(index + bar_width)
    ax.set_xticklabels(bandwidths)
    ax.legend()

    # 使用 ScalarFormatter 来设置 y 轴为科学计数法，并显示 e2 单位
    formatter = ticker.ScalarFormatter(useMathText=True)
    formatter.set_powerlimits((0, 0))  # 强制使用科学计数法
    ax.yaxis.set_major_formatter(formatter)

    # 在 y 轴上方显示 'e2' 单位
    ax.yaxis.get_offset_text().set_visible(False)  # 隐藏默认的偏移量标记
    ax.annotate('e2', xy=(0, 1), xytext=(10, 2), xycoords='axes fraction', textcoords='offset points',
                ha='center', va='baseline', fontsize=12)
    ax.set_ylim(0, ylim)  # 手动设置 y 轴的范围

# 绘制 VGG 的柱状图
plot_bars(axs[0], baseline_vgg, topk_vgg, netsense_vgg, '(a) VGG', 'Speed (images/sec)', 300)

# 绘制 ResNet 的柱状图
plot_bars(axs[1], baseline_resnet, topk_resnet, netsense_resnet, '(b) ResNet', 'Speed (images/sec)', 300)

# 绘制 BERT 的柱状图
plot_bars(axs[2], baseline_bert, topk_bert, netsense_bert, '(c) BERT', 'Speed (questions/sec)', 100)

# 调整布局并显示图像
plt.tight_layout()
plt.show()
plt.savefig('training_throughput_s1.pdf', format='pdf', bbox_inches='tight')
